"""
Filename: ImprovedSENet.py
Author: Deng Weiwei
Description: This script contains improvements to SENet by replacing the Linear layer in the SE block with a 1x1 convolution.
"""
import torch
import torch.nn as nn
import torch.nn.functional as F
from torchsummary import summary

class MHSA(nn.Module):
    def __init__(self, n_dims, width, height, heads=4):
        super(MHSA, self).__init__()
        self.heads = heads

        self.query = nn.Conv2d(n_dims, n_dims, kernel_size=1)
        self.key = nn.Conv2d(n_dims, n_dims, kernel_size=1)
        self.value = nn.Conv2d(n_dims, n_dims, kernel_size=1)

        self.rel_h = nn.Parameter(torch.randn([1, heads, n_dims // heads, 1, height]), requires_grad=True)
        self.rel_w = nn.Parameter(torch.randn([1, heads, n_dims // heads, width, 1]), requires_grad=True)

        self.softmax = nn.Softmax(dim=-1)

    def forward(self, x):
        n_batch, C, width, height = x.size()
        q = self.query(x).view(n_batch, self.heads, C // self.heads, -1)
        k = self.key(x).view(n_batch, self.heads, C // self.heads, -1)
        v = self.value(x).view(n_batch, self.heads, C // self.heads, -1)

        content_content = torch.matmul(q.permute(0, 1, 3, 2), k)

        content_position = (self.rel_h + self.rel_w).view(1, self.heads, C // self.heads, -1).permute(0, 1, 3, 2)
        content_position = torch.matmul(content_position, q)

        energy = content_content + content_position
        attention = self.softmax(energy)

        out = torch.matmul(v, attention.permute(0, 1, 3, 2))
        out = out.view(n_batch, C, width, height)

        return out
    
class Bottleneck(nn.Module):
    expansion = 4

    def __init__(self, in_channels, mid_channels, stride=1, heads=4, mhsa=False, resolution=None,add_se=False):
        super(Bottleneck, self).__init__()

        self.conv1 = nn.Conv2d(in_channels, mid_channels, kernel_size=1, bias=False)
        self.bn1 = nn.BatchNorm2d(mid_channels)
        
        if not mhsa:
            self.conv2 = nn.Conv2d(mid_channels, mid_channels, kernel_size=3, padding=1, stride=stride, bias=False)
        else:
            self.conv2 = nn.ModuleList()
            self.conv2.append(MHSA(mid_channels, width=int(resolution[0]), height=int(resolution[1]), heads=heads))
            if stride == 2:
                self.conv2.append(nn.AvgPool2d(2, 2))
            self.conv2 = nn.Sequential(*self.conv2)
        self.bn2 = nn.BatchNorm2d(mid_channels)
        
        self.conv3 = nn.Conv2d(mid_channels, self.expansion * mid_channels, kernel_size=1, bias=False)
        self.bn3 = nn.BatchNorm2d(self.expansion * mid_channels)
        
        self.residual  = nn.Sequential()
        if stride != 1 or in_channels != self.expansion*mid_channels:
            self.residual  = nn.Sequential(
                nn.Conv2d(in_channels, self.expansion*mid_channels, kernel_size=1, stride=stride),
                nn.BatchNorm2d(self.expansion*mid_channels)
            )
            
        
        if add_se:
            # Add SE module after the last convolution layer
            self.se = nn.Sequential(
                nn.AdaptiveAvgPool2d(1),
                nn.Conv2d(self.expansion * mid_channels, self.expansion * mid_channels // 4, kernel_size=1),
                nn.ReLU(inplace=True),
                nn.Conv2d(self.expansion * mid_channels // 4, self.expansion * mid_channels, kernel_size=1),
                nn.Sigmoid()
            )
        else:
            self.se = None

    def forward(self, x):
        out = F.relu(self.bn1(self.conv1(x)))
        out = F.relu(self.bn2(self.conv2(out)))
        out = self.bn3(self.conv3(out))
        residual = self.residual (x)
        if self.se is not None:
            # 计算SE模块的权重
            se_weight = self.se(out)
            # 对特征图应用SE权重
            out = out * se_weight
        
        out += residual
        out = F.relu(out)
        return out
    
class ResNet(nn.Module):
    def __init__(self, block, num_blocks, num_classes, resolution, heads):
        super(ResNet, self).__init__()
        self.in_channels = 64
        self.resolution = list(resolution)

        self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, bias=False)
        self.resolution[0] /= 2
        self.resolution[1] /= 2
        
        self.bn1 = nn.BatchNorm2d(64)
        self.relu = nn.ReLU(inplace=True)
        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
        self.resolution[0] /= 2
        self.resolution[1] /= 2

        self.layer1 = self._make_layer(block, 64, num_blocks[0], stride=1,add_se=True)
        self.layer2 = self._make_layer(block, 128, num_blocks[1], stride=2,add_se=True)
        self.layer3 = self._make_layer(block, 256, num_blocks[2], stride=2, add_se=True)
        self.layer4 = self._make_layer(block, 512, num_blocks[3], stride=2, add_se=True)

        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
        self.fc = nn.Sequential(
            nn.Dropout(0.3), 
            nn.Linear(512 * block.expansion, num_classes)
        )

    def _make_layer(self, block, mid_channels, num_blocks, stride=1, heads=4, mhsa=False,add_se=False):
        out_channels = mid_channels * block.expansion
        layers = []
        
        layers.append(block(self.in_channels, mid_channels, stride, heads, mhsa, self.resolution.copy(),add_se))
        
        if stride == 2:
            self.resolution[0] /= 2
            self.resolution[1] /= 2
            
        self.in_channels = out_channels
        for _ in range(1, num_blocks):
            layers.append(block(self.in_channels, mid_channels,1, heads, mhsa, self.resolution.copy(),add_se))
        return nn.Sequential(*layers)

    def forward(self, x):
        out = self.relu(self.bn1(self.conv1(x)))
        out = self.maxpool(out)

        out = self.layer1(out)
        out = self.layer2(out)
        out = self.layer3(out)
        out = self.layer4(out)

        out = self.avgpool(out)
        out = torch.flatten(out, 1)
        out = self.fc(out)
        return out

def SENet(num_classes=100, resolution=(32,32), heads=16):
    return ResNet(Bottleneck, [2, 2, 2, 2], num_classes=num_classes, resolution=resolution, heads=heads)

if __name__ == "__main__":
    net = SENet()
    summary(net, input_size=(3, 32, 32), device='cuda')